import re
import cv2
import PIL.Image
import numpy as np


def draw_frame(image, frame_all, label=None, save_path=None):
    frame_all = frame_all.copy().astype(np.int32)
    if len(frame_all.shape) == 1:
        frame_all = np.array([frame_all])
    center = frame_all[:, 0:2].copy()
    frame_all[:, 0:2] = frame_all[:, 0:2] - frame_all[:, 2:4] / 2
    frame_all[:, 2:4] = frame_all[:, 0:2] + frame_all[:, 2:4]  # (x_min,y_min,x_max,y_max)
    for index, frame in enumerate(frame_all):
        x1, y1, x2, y2 = frame[0:4]
        cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color=(0, 255, 0), thickness=1)
        label_ = label[index] if label is not None else index
        cx, cy = center[index]
        cv2.putText(image, f'{label_}', org=(cx, cy), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4,
                    color=(0, 255, 0), thickness=1)
    image = PIL.Image.fromarray(image)
    if save_path is None:
        image.show()
    else:
        image.save(save_path)


with open('train.txt', 'r', encoding='utf-8') as f:
    line_all = f.readlines()
line = line_all[0]
regex = re.compile(r'(.*\..{3,4}) (.*)')
search = regex.search(line.strip())
image_path = search.group(1)
label_path = search.group(2)
image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)  # 读取图片
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转为RGB通道
with open(label_path, 'r', encoding='utf-8') as f:
    line_all = f.readlines()
label = []
for line in line_all:
    cls, x, y, w, h = line.strip().split()
    label.append([x, y, w, h, cls])
label = np.array(label, dtype=np.float64)
h, w, _ = image.shape
label[:, 0] *= w
label[:, 1] *= h
label[:, 2] *= w
label[:, 3] *= h
draw_frame(image, label, label[:, 4])
